
import os, bpy, bpy.ops, math, mathutils, ossl.util as util
from bpy.types import Operator
from mathutils import Vector
from math import pi

class Spring:
    
    def generate_default():
        util.setMode("OBJECT")

        revolutions = 8
        stretch_factor = 6
        spring_thickness = 0.1
        vertex_ring_count = 16 # must be even value

        #### create new scene
        scene = bpy.data.scenes.new(name="SpringScene")
        bpy.context.window.scene = scene

        #### create spring
        bpy.ops.curve.spirals(spiral_type='ARCH', dif_z=spring_thickness * stretch_factor, turns=revolutions)
        bpy.context.object.data.bevel_depth = spring_thickness / 2
        bpy.context.object.data.bevel_resolution = int((vertex_ring_count - 4) / 2) # #of vertices per ring = 4 + res*2
        util.setMode("OBJECT")
        bpy.ops.object.convert(target='MESH')
        spiral = bpy.context.active_object
        height = spiral.dimensions.z

        #### create faces (start and end)
        def addFace(start):
            for i in range(start, start + vertex_ring_count): spiral.data.vertices[i].select = True
            util.setMode("EDIT")
            bpy.ops.mesh.edge_face_add()
            util.setMode("OBJECT")
            for i in range(start, start + vertex_ring_count): spiral.data.vertices[i].select = False
        addFace(0)
        addFace(len(spiral.data.vertices) - vertex_ring_count)

        #### create armature
        util.setMode("OBJECT")
        bpy.ops.object.armature_add(enter_editmode=False, align='WORLD', location=(0, 0, 0), scale=(1,1,1))

        #### @Ada : hinge joint vs. ball joint vs. carpal joint

        #### create mPelvis
        util.setMode("EDIT")
        arm = bpy.context.active_object
        pelvis = arm.data.edit_bones["Bone"]
        pelvis.name = "mPelvis"
        pelvis.length = 0.2
        pelvis.head.z = 0 # height/2
        pelvis.tail.z = 0.2 # (height/2) + 0.2

        #### create mSpine1
        bpy.ops.armature.extrude_move(
            ARMATURE_OT_extrude={"forked":False}, TRANSFORM_OT_translate={"value":(0, 0, 0.2),
            "orient_type":'GLOBAL', "orient_matrix":((1, 0, 0), (0, 1, 0), (0, 0, 1)), "orient_matrix_type":'GLOBAL'})
        spine = arm.data.edit_bones[pelvis.name + ".001"]
        spine.name = "mSpine1"
        spine.parent = None # to release the "connected" state and the dependant height of the bone's head
        spine.parent = pelvis
        spine.head.z = height
        spine.tail.z = height + 0.2

        #### parent spring to armature
        util.setMode("OBJECT")
        spiral.select_set(True)
        bpy.context.view_layer.objects.active = arm
        bpy.ops.object.parent_set(type="ARMATURE_NAME") # vertex groups exist, but empty

        #### edit spring for collecting spring vertices
        bpy.context.view_layer.objects.active = None
        arm.select_set(False)
        spiral.select_set(False)
        spiral.select_set(True)
        bpy.context.view_layer.objects.active = spiral
        util.setMode("EDIT")

        #### collect spring vertex rings
        def addVerts(i, sv):
            zmin = sv[i].co.z
            zmax = zmin
            for k in range(i+1, i+vertex_ring_count):
                if sv[k].co.z < zmin: zmin = sv[k].co.z
                if sv[k].co.z > zmax: zmax = sv[k].co.z
            z = zmin + (zmax-zmin)/2 # median point (z) per vertex ring
            svg.append((int(i/vertex_ring_count), z, range(i, i+vertex_ring_count)))

        bpy.ops.mesh.select_all(action='SELECT')
        sel_verts = spiral.data.vertices # vertex rings are ordered z-down-top by creation 
        svg = [] # for vertex rings, 12 vertices each ring by creation
        for i in range(0, int(len(sel_verts)/vertex_ring_count)): addVerts(int(i*vertex_ring_count), sel_verts)
        zmin = svg[0][1]
        zmax = zmin # find lower and upper z of spring
        for i in range(1, len(svg)):
            if svg[i][1] < zmin: zmin = svg[i][1]
            if svg[i][1] > zmax: zmax = svg[i][1]

        #### weight painting
        def setWeight(i, f):
            vg = spiral.vertex_groups[i]
            def ff(f, z): # scale vertex ring height to weight
                zd = (z-zmin)/(zmax-zmin)
                if f == 1: zd = 1-zd
                return zd
            for (i, z, vs) in svg: vg.add(vs, ff(f, z), 'REPLACE')
        util.setMode("WEIGHT_PAINT")
        setWeight(0,1) # 0: vertex group "mPelvis"
        setWeight(1,-1) # 1: vertex group "mSpine1"

        #### create animation

        def selectBone(name):
            bone = arm.data.bones[name]
            bone.select = True
            arm.data.bones.active = bone
            return bone

        def selectPoseBone(name):
            selectBone(name)
            return arm.pose.bones[name]

        def unselectBones():
            bpy.ops.pose.select_all(action='DESELECT')

        def insertRestPose():
            unselectBones()
            selectPoseBone("mPelvis")
            selectPoseBone("mSpine1")
            bpy.ops.anim.keyframe_insert_menu(type='Location')
            #bpy.ops.anim.keyframe_insert_menu(type='Rotation')

        def insertLoopPose():
            unselectBones()
            mSpine = selectPoseBone("mSpine1")
            h0 = revolutions * spring_thickness 
            mSpine.location = Vector((0,h0-height, 0)) # minimum height of mSpine1
            bpy.ops.anim.keyframe_insert_menu(type='Location')
            #bpy.ops.anim.keyframe_insert_menu(type='Rotation')

        util.setMode("OBJECT")
        spiral.select_set(False)
        arm.select_set(True)
        bpy.context.view_layer.objects.active = arm
        util.setMode("POSE")

        bpy.context.scene.frame_set(0)
        insertRestPose()

        bpy.context.scene.frame_set(64)
        insertRestPose()

        bpy.context.scene.frame_set(32)
        insertLoopPose()

        ## looping and total frames
        bpy.context.scene.frame_set(1)
        bpy.context.scene.timeline_markers.new("loop_first", frame=1)

        bpy.context.scene.frame_set(63)
        bpy.context.scene.timeline_markers.new("loop_last", frame=63)

        bpy.context.scene.frame_set(0)
        bpy.context.scene.frame_start = 0
        bpy.context.scene.frame_end = 64

        #### standardizing spring (mesh object) orientation
        util.setMode("OBJECT")
        spiral.select_set(True)
        arm.select_set(True)
        bpy.context.view_layer.objects.active = arm
        bpy.ops.transform.rotate(value=-pi/4, orient_axis='Z', orient_type='GLOBAL')
        bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)

        #### automatic export of animation using OS/SL plugin
        util.setMode("OBJECT")
        from ossl.types.Animation.export import export_anim
        print("start anim output")
        bpy.ops.ossl.export_animation_anim("INVOKE_DEFAULT")
        print("end anim output")

        #### automatic export of Collada not possible here
        #### because of non-modality of both exports
        #### hitherto calling Collada-export overrides call of anim-export
        #### do it in standard Collada export, using preset "rigged" plus checking "joints"
        #util.setMode("OBJECT")
        #arm.select_set(True)
        #spiral.select_set(True)
        #bpy.context.view_layer.objects.active = spiral
        #bpy.ops.wm.collada_export("INVOKE_DEFAULT")

        print("*** ok ***")

class OSSL_OT_Spring_generate_Default(bpy.types.Operator):
    """Generate Default OSSL Spring with Test Animation"""
    bl_label = "OSSL Menu Pool.Spring Generate Default"
    bl_idname = "ossl.spring_generate_default"
    bl_options = { 'REGISTER' }

    def draw(self, context):
        layout = self.layout

    @classmethod
    def poll(cls, context):
        return True
    
    def execute(self, context):
        print("dummy Spring generation")
        Spring.generate_default()
        print("dummy Spring generation done")
        return {'FINISHED'}

def OSSL_operator_Spring_generate_Default(self, comtext):
    self.layout.operator(OSSL_OT_Spring_generate_Default.bl_idname, text="generate default", icon='PREFERENCES')

import ossl.util as util

def register():
    print("registering: OSSL Type Pool.Spring Generate Default")
    bpy.utils.register_class(OSSL_OT_Spring_generate_Default)
    bpy.types.OSSL_MT_Spring.append(OSSL_operator_Spring_generate_Default)
    print("registered: OSSL Type Pool.Spring Generate Default")

def unregister():
    print("unregistering: OSSL Type Pool.Spring Generate Default")
    bpy.types.OSSL_MT_Spring.remove(OSSL_operator_Spring_generate_Default)
    bpy.utils.unregister_class(OSSL_OT_Spring_generate_Default)
    print("unregistered: OSSL Type Pool.Spring Generate Default")

if __name__ == "__main__":
    import ossl.util as util
    util.reset_addon()
